/*
 * 3D Helmholtz Finite Difference Method Solver
 * Solves: (I - κ∇²)u = f in 3D domain
 * Where I is identity matrix, κ is Helmholtz parameter, ∇² is Laplacian
 * 
 * The matrix is constructed to be diagonally dominant with diagonal elements near 1
 * Uses 7-point stencil for 3D finite differences
 * 
 * Now supports half single precision (16-bit floating point)
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <mpi.h>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

// Half precision floating point support
#ifdef __GNUC__
    // GCC supports _Float16
    #if __GNUC__ >= 6
        typedef _Float16 half_float;
        #define HALF_PRECISION_SUPPORTED 1
        #define sinf16(x) ((half_float)sinf((float)(x)))
        #define cosf16(x) ((half_float)cosf((float)(x)))
        #define sqrtf16(x) ((half_float)sqrtf((float)(x)))
        #define fabsf16(x) ((half_float)fabsf((float)(x)))
        #define fminf16(x,y) ((half_float)fminf((float)(x), (float)(y)))
        #define fmaxf16(x,y) ((half_float)fmaxf((float)(x), (float)(y)))
    #else
        typedef float half_float;
        #define HALF_PRECISION_SUPPORTED 0
        #define sinf16(x) sinf(x)
        #define cosf16(x) cosf(x)
        #define sqrtf16(x) sqrtf(x)
        #define fabsf16(x) fabsf(x)
        #define fminf16(x,y) fminf(x,y)
        #define fmaxf16(x,y) fmaxf(x,y)
    #endif
#elif defined(__clang__)
    // Clang supports __fp16
    #if __has_extension(c_half_precision)
        typedef __fp16 half_float;
        #define HALF_PRECISION_SUPPORTED 1
        #define sinf16(x) ((__fp16)sinf((float)(x)))
        #define cosf16(x) ((__fp16)cosf((float)(x)))
        #define sqrtf16(x) ((__fp16)sqrtf((float)(x)))
        #define fabsf16(x) ((__fp16)fabsf((float)(x)))
        #define fminf16(x,y) ((__fp16)fminf((float)(x), (float)(y)))
        #define fmaxf16(x,y) ((__fp16)fmaxf((float)(x), (float)(y)))
    #else
        typedef float half_float;
        #define HALF_PRECISION_SUPPORTED 0
        #define sinf16(x) sinf(x)
        #define cosf16(x) cosf(x)
        #define sqrtf16(x) sqrtf(x)
        #define fabsf16(x) fabsf(x)
        #define fminf16(x,y) fminf(x,y)
        #define fmaxf16(x,y) fmaxf(x,y)
    #endif
#else
    // Fallback to single precision float
    typedef float half_float;
    #define HALF_PRECISION_SUPPORTED 0
    #define sinf16(x) sinf(x)
    #define cosf16(x) cosf(x)
    #define sqrtf16(x) sqrtf(x)
    #define fabsf16(x) fabsf(x)
    #define fminf16(x,y) fminf(x,y)
    #define fmaxf16(x,y) fmaxf(x,y)
#endif

typedef struct {
    int nx, ny, nz;      // Grid dimensions
    half_float dx, dy, dz;   // Grid spacing
    half_float kappa;        // Helmholtz parameter
    half_float *u;           // Solution vector
    half_float *f;           // Right-hand side
    half_float *matrix;      // Coefficient matrix (7-diagonal storage)
    half_float *work;        // Working arrays
} HelmholtzFDM3D;

// Matrix storage indices for 7-point stencil
enum {
    CENTER = 0,    // (i,j,k)
    WEST = 1,      // (i-1,j,k)  
    EAST = 2,      // (i+1,j,k)
    SOUTH = 3,     // (i,j-1,k)
    NORTH = 4,     // (i,j+1,k)
    BOTTOM = 5,    // (i,j,k-1)
    TOP = 6        // (i,j,k+1)
};

// Convert 3D indices to linear index
static inline int idx3d(int i, int j, int k, int nx, int ny, int nz) {
    return i + nx * (j + ny * k);
}

// Initialize the Helmholtz FDM solver
HelmholtzFDM3D* helmholtz_fdm_init(int nx, int ny, int nz, 
                                   half_float lx, half_float ly, half_float lz,
                                   half_float kappa) {
    HelmholtzFDM3D* solver = malloc(sizeof(HelmholtzFDM3D));
    
    solver->nx = nx;
    solver->ny = ny; 
    solver->nz = nz;
    solver->dx = lx / (half_float)(nx - 1);
    solver->dy = ly / (half_float)(ny - 1);
    solver->dz = lz / (half_float)(nz - 1);
    solver->kappa = kappa;
    
    int total_points = nx * ny * nz;
    
    // Allocate memory
    solver->u = calloc(total_points, sizeof(half_float));
    solver->f = calloc(total_points, sizeof(half_float));
    solver->matrix = calloc(7 * total_points, sizeof(half_float));
    solver->work = calloc(total_points, sizeof(half_float));
    
    if (!solver->u || !solver->f || !solver->matrix || !solver->work) {
        fprintf(stderr, "Memory allocation failed\n");
        exit(1);
    }
    
    return solver;
}

// Set up the coefficient matrix for (I - κ∇²)u = f
void helmholtz_fdm_setup_matrix(HelmholtzFDM3D* solver) {
    int nx = solver->nx, ny = solver->ny, nz = solver->nz;
    half_float dx = solver->dx, dy = solver->dy, dz = solver->dz;
    half_float kappa = solver->kappa;
    
    // Finite difference coefficients for second derivatives
    half_float cx = kappa / (dx * dx);
    half_float cy = kappa / (dy * dy);
    half_float cz = kappa / (dz * dz);
    
    printf("Setting up matrix with kappa=%.6f, dx=%.6f, dy=%.6f, dz=%.6f\n", 
           (float)kappa, (float)dx, (float)dy, (float)dz);
    printf("Coefficients: cx=%.6f, cy=%.6f, cz=%.6f\n", (float)cx, (float)cy, (float)cz);
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                
                // Initialize all stencil points to zero
                for (int m = 0; m < 7; m++) {
                    solver->matrix[7 * idx + m] = (half_float)0.0;
                }
                
                // Interior points use full 7-point stencil
                if (i > 0 && i < nx-1 && j > 0 && j < ny-1 && k > 0 && k < nz-1) {
                    // Center point: 1 + κ(2/dx² + 2/dy² + 2/dz²) for (I - κ∇²)
                    solver->matrix[7 * idx + CENTER] = (half_float)1.0 + (half_float)2.0 * (cx + cy + cz);
                    
                    // Neighbor points: -κ/h²
                    solver->matrix[7 * idx + WEST]   = -cx;
                    solver->matrix[7 * idx + EAST]   = -cx;
                    solver->matrix[7 * idx + SOUTH]  = -cy;
                    solver->matrix[7 * idx + NORTH]  = -cy;
                    solver->matrix[7 * idx + BOTTOM] = -cz;
                    solver->matrix[7 * idx + TOP]    = -cz;
                }
                // Boundary points: Homogeneous Dirichlet (u = 0)
                else {
                    solver->matrix[7 * idx + CENTER] = (half_float)1.0;
                    // All other coefficients remain zero for boundary
                }
            }
        }
    }
    
    printf("Matrix setup complete. Sample diagonal element: %.6f\n", 
           (float)solver->matrix[7 * idx3d(nx/2, ny/2, nz/2, nx, ny, nz) + CENTER]);
}

// Set right-hand side with analytical test function
void helmholtz_fdm_set_rhs(HelmholtzFDM3D* solver) {
    int nx = solver->nx, ny = solver->ny, nz = solver->nz;
    half_float dx = solver->dx, dy = solver->dy, dz = solver->dz;
    half_float kappa = solver->kappa;
    
    // Test function: u(x,y,z) = sin(πx/Lx) * sin(πy/Ly) * sin(πz/Lz)
    // For this function, -∇²u = π²(1/Lx² + 1/Ly² + 1/Lz²) * u
    half_float lx = (half_float)(nx - 1) * dx;
    half_float ly = (half_float)(ny - 1) * dy;
    half_float lz = (half_float)(nz - 1) * dz;
    
    half_float lambda = (half_float)(M_PI * M_PI) * ((half_float)1.0/(lx*lx) + (half_float)1.0/(ly*ly) + (half_float)1.0/(lz*lz));
    
    printf("Setting RHS with lambda=%.6f, domain size: %.3f x %.3f x %.3f\n", 
           (float)lambda, (float)lx, (float)ly, (float)lz);
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                
                half_float x = (half_float)i * dx;
                half_float y = (half_float)j * dy; 
                half_float z = (half_float)k * dz;
                
                // u(x,y,z) = sin(πx/Lx) * sin(πy/Ly) * sin(πz/Lz)
                half_float u_exact = sinf16((half_float)M_PI * x / lx) * sinf16((half_float)M_PI * y / ly) * sinf16((half_float)M_PI * z / lz);
                
                // (I - κ∇²)u = u + κλu = (1 + κλ)u
                if (i == 0 || i == nx-1 || j == 0 || j == ny-1 || k == 0 || k == nz-1) {
                    // Boundary: u = 0
                    solver->f[idx] = (half_float)0.0;
                } else {
                    // Interior: f = (1 + κλ)u
                    solver->f[idx] = ((half_float)1.0 + kappa * lambda) * u_exact;
                }
            }
        }
    }
    
    printf("RHS setup complete. Max RHS value: %.6f\n", 
           (float)solver->f[idx3d(nx/2, ny/2, nz/2, nx, ny, nz)]);
}

// Matrix-vector multiplication: y = A*x
void helmholtz_fdm_matvec(HelmholtzFDM3D* solver, half_float* x, half_float* y) {
    int nx = solver->nx, ny = solver->ny, nz = solver->nz;
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                half_float* coeff = &solver->matrix[7 * idx];
                
                y[idx] = coeff[CENTER] * x[idx];
                
                // Add neighbor contributions if they exist
                if (i > 0)    y[idx] += coeff[WEST]   * x[idx3d(i-1, j, k, nx, ny, nz)];
                if (i < nx-1) y[idx] += coeff[EAST]   * x[idx3d(i+1, j, k, nx, ny, nz)];
                if (j > 0)    y[idx] += coeff[SOUTH]  * x[idx3d(i, j-1, k, nx, ny, nz)];
                if (j < ny-1) y[idx] += coeff[NORTH]  * x[idx3d(i, j+1, k, nx, ny, nz)];
                if (k > 0)    y[idx] += coeff[BOTTOM] * x[idx3d(i, j, k-1, nx, ny, nz)];
                if (k < nz-1) y[idx] += coeff[TOP]    * x[idx3d(i, j, k+1, nx, ny, nz)];
            }
        }
    }
}

// Compute vector 2-norm
half_float vector_norm(half_float* v, int n) {
    half_float sum = (half_float)0.0;
    for (int i = 0; i < n; i++) {
        sum += v[i] * v[i];
    }
    return sqrtf16(sum);
}

// Simple Jacobi iterative solver
int helmholtz_fdm_solve_jacobi(HelmholtzFDM3D* solver, half_float tol, int max_iter) {
    int n = solver->nx * solver->ny * solver->nz;
    half_float* u_old = solver->work;
    
    printf("Starting Jacobi solver with tol=%.2e, max_iter=%d\n", (float)tol, max_iter);
    
    // Initialize solution to zero
    memset(solver->u, 0, n * sizeof(half_float));
    
    half_float final_residual = (half_float)1e10;
    
    for (int iter = 0; iter < max_iter; iter++) {
        // Copy current solution
        memcpy(u_old, solver->u, n * sizeof(half_float));
        
        // Jacobi update: u_new[i] = (f[i] - A_off[i]*u_old) / A_diag[i]
        for (int k = 0; k < solver->nz; k++) {
            for (int j = 0; j < solver->ny; j++) {
                for (int i = 0; i < solver->nx; i++) {
                    int idx = idx3d(i, j, k, solver->nx, solver->ny, solver->nz);
                    half_float* coeff = &solver->matrix[7 * idx];
                    
                    half_float rhs = solver->f[idx];
                    
                    // Subtract off-diagonal terms
                    if (i > 0)           rhs -= coeff[WEST]   * u_old[idx3d(i-1, j, k, solver->nx, solver->ny, solver->nz)];
                    if (i < solver->nx-1) rhs -= coeff[EAST]   * u_old[idx3d(i+1, j, k, solver->nx, solver->ny, solver->nz)];
                    if (j > 0)           rhs -= coeff[SOUTH]  * u_old[idx3d(i, j-1, k, solver->nx, solver->ny, solver->nz)];
                    if (j < solver->ny-1) rhs -= coeff[NORTH]  * u_old[idx3d(i, j+1, k, solver->nx, solver->ny, solver->nz)];
                    if (k > 0)           rhs -= coeff[BOTTOM] * u_old[idx3d(i, j, k-1, solver->nx, solver->ny, solver->nz)];
                    if (k < solver->nz-1) rhs -= coeff[TOP]    * u_old[idx3d(i, j, k+1, solver->nx, solver->ny, solver->nz)];
                    
                    // Update solution
                    solver->u[idx] = rhs / coeff[CENTER];
                }
            }
        }
        
        // Check convergence every 10 iterations
        if (iter % 10 == 0) {
            // Compute residual: r = f - A*u
            helmholtz_fdm_matvec(solver, solver->u, solver->work);
            for (int i = 0; i < n; i++) {
                solver->work[i] = solver->f[i] - solver->work[i];
            }
            
            half_float residual_norm = vector_norm(solver->work, n);
            final_residual = residual_norm;
            printf("Iteration %d: residual norm = %.6e\n", iter, (float)residual_norm);
            
            if (residual_norm < tol) {
                printf("Converged in %d iterations\n", iter);
                return iter;
            }
        }
    }
    
    printf("Maximum iterations reached. Final residual: %.6e\n", (float)final_residual);
    return max_iter;
}

// Compute analytical solution for verification
void helmholtz_fdm_analytical_solution(HelmholtzFDM3D* solver, half_float* u_exact) {
    int nx = solver->nx, ny = solver->ny, nz = solver->nz;
    half_float dx = solver->dx, dy = solver->dy, dz = solver->dz;
    
    half_float lx = (half_float)(nx - 1) * dx;
    half_float ly = (half_float)(ny - 1) * dy;
    half_float lz = (half_float)(nz - 1) * dz;
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                
                half_float x = (half_float)i * dx;
                half_float y = (half_float)j * dy;
                half_float z = (half_float)k * dz;
                
                // Boundary points have u = 0
                if (i == 0 || i == nx-1 || j == 0 || j == ny-1 || k == 0 || k == nz-1) {
                    u_exact[idx] = (half_float)0.0;
                } else {
                    u_exact[idx] = sinf16((half_float)M_PI * x / lx) * sinf16((half_float)M_PI * y / ly) * sinf16((half_float)M_PI * z / lz);
                }
            }
        }
    }
}

// Cleanup memory
void helmholtz_fdm_destroy(HelmholtzFDM3D* solver) {
    if (solver) {
        free(solver->u);
        free(solver->f);
        free(solver->matrix);
        free(solver->work);
        free(solver);
    }
}

// Test function
int test_helmholtz_fdm_3d() {
    printf("=== 3D Helmholtz FDM Solver Test (Half Precision) ===\n");
    
    #if HALF_PRECISION_SUPPORTED
        printf("Using true 16-bit half precision floating point\n");
    #else
        printf("Using single precision floating point (fallback)\n");
    #endif
    
    // Problem parameters
    int nx = 32, ny = 32, nz = 32;
    half_float lx = (half_float)1.0, ly = (half_float)1.0, lz = (half_float)1.0;
    half_float kappa = (half_float)0.01;  // Smaller value for better diagonal dominance
    
    printf("Grid: %dx%dx%d, Domain: %.1fx%.1fx%.1f, kappa=%.3f\n", 
           nx, ny, nz, (float)lx, (float)ly, (float)lz, (float)kappa);
    
    // Initialize solver
    HelmholtzFDM3D* solver = helmholtz_fdm_init(nx, ny, nz, lx, ly, lz, kappa);
    
    // Setup matrix and RHS
    helmholtz_fdm_setup_matrix(solver);
    helmholtz_fdm_set_rhs(solver);
    
    // Solve the system
    int iterations = helmholtz_fdm_solve_jacobi(solver, (half_float)1e-6, 1000);  // Relaxed tolerance for half precision
    
    // Verify against analytical solution
    half_float* u_exact = malloc(nx * ny * nz * sizeof(half_float));
    helmholtz_fdm_analytical_solution(solver, u_exact);
    
    // Compute error
    half_float error = (half_float)0.0;
    for (int i = 0; i < nx * ny * nz; i++) {
        half_float diff = solver->u[i] - u_exact[i];
        error += diff * diff;
    }
    error = sqrtf16(error) / sqrtf16((half_float)(nx * ny * nz));
    
    printf("L2 error vs analytical solution: %.6e\n", (float)error);
    printf("Solution norm: %.6e\n", (float)vector_norm(solver->u, nx * ny * nz));
    printf("Exact solution norm: %.6e\n", (float)vector_norm(u_exact, nx * ny * nz));
    printf("Iterations used: %d\n", iterations);
    
    // Check matrix diagonal dominance
    int dominant_count = 0;
    half_float min_diag = (half_float)1e10, max_diag = (half_float)(-1e10);
    for (int k = 1; k < nz-1; k++) {
        for (int j = 1; j < ny-1; j++) {
            for (int i = 1; i < nx-1; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                half_float* coeff = &solver->matrix[7 * idx];
                
                half_float diag = fabsf16(coeff[CENTER]);
                half_float off_diag_sum = fabsf16(coeff[WEST]) + fabsf16(coeff[EAST]) + 
                                         fabsf16(coeff[SOUTH]) + fabsf16(coeff[NORTH]) + 
                                         fabsf16(coeff[BOTTOM]) + fabsf16(coeff[TOP]);
                
                if (diag > off_diag_sum) dominant_count++;
                
                min_diag = fminf16(min_diag, coeff[CENTER]);
                max_diag = fmaxf16(max_diag, coeff[CENTER]);
            }
        }
    }
    
    int total_interior = (nx-2) * (ny-2) * (nz-2);
    printf("Diagonally dominant points: %d/%d (%.1f%%)\n", 
           dominant_count, total_interior, 100.0 * dominant_count / total_interior);
    printf("Diagonal element range: [%.6f, %.6f]\n", (float)min_diag, (float)max_diag);
    
    // Success criteria: relaxed error tolerance for half precision
    half_float tolerance = (half_float)1e-2;  // More relaxed for half precision
    int success = (error < tolerance);
    printf("Test result: %s (error=%.2e < %.2e)\n", success ? "PASSED" : "FAILED", (float)error, (float)tolerance);
    
    free(u_exact);
    helmholtz_fdm_destroy(solver);
    
    return success ? 0 : 1;  // Return 0 for success
}

// Main function for standalone testing
#ifdef HELMHOLTZ_FDM_MAIN
int main(int argc, char** argv) {
    return test_helmholtz_fdm_3d();
}
#endif 